# ======================================================================
# Copyright TOTAL / CERFACS / LIRMM (03/2020)
# Contributor: Adrien Suau (<adrien.suau@cerfacs.fr>
#                           <adrien.suau@lirmm.fr>)
#
# This software is governed by the CeCILL-B license under French law and
# abiding  by the  rules of  distribution of free software. You can use,
# modify  and/or  redistribute  the  software  under  the  terms  of the
# CeCILL-B license as circulated by CEA, CNRS and INRIA at the following
# URL "http://www.cecill.info".
#
# As a counterpart to the access to  the source code and rights to copy,
# modify and  redistribute granted  by the  license, users  are provided
# only with a limited warranty and  the software's author, the holder of
# the economic rights,  and the  successive licensors  have only limited
# liability.
#
# In this respect, the user's attention is drawn to the risks associated
# with loading,  using, modifying and/or  developing or reproducing  the
# software by the user in light of its specific status of free software,
# that  may mean  that it  is complicated  to manipulate,  and that also
# therefore  means that  it is reserved for  developers and  experienced
# professionals having in-depth  computer knowledge. Users are therefore
# encouraged  to load and  test  the software's  suitability as  regards
# their  requirements  in  conditions  enabling  the  security  of their
# systems  and/or  data to be  ensured and,  more generally,  to use and
# operate it in the same conditions as regards security.
#
# The fact that you  are presently reading this  means that you have had
# knowledge of the CeCILL-B license and that you accept its terms.
# ======================================================================

import typing as ty
from concurrent.futures import ProcessPoolExecutor
import multiprocessing
import os
from itertools import product
import logging

from qaths._cli.data_generation.utils.ranges import RangeLike

# See
# https://stackoverflow.com/questions/15639779/why-does-multiprocessing-use-only-a-single-core-after-i-import-numpy
# Provided snippet does not seems to work correctly, so I found another solution
_cpu_count = multiprocessing.cpu_count()
_all_cpus = list(range(_cpu_count))
os.sched_setaffinity(os.getpid(), _all_cpus)

logger = logging.getLogger("qaths._cli.data_generation.utils.parallel")


def execute_task_in_parallel(
    task: ty.Callable, ranges: ty.List[RangeLike], verbose: bool = True,
) -> ty.List[ty.Any]:
    total_task_number = 1
    for range in ranges:
        total_task_number *= len(range)
    results = []
    with ProcessPoolExecutor(max_workers=_cpu_count) as executor:
        for idx, result in enumerate(executor.map(task, product(*ranges))):
            results.append(result)
            if verbose:
                print(f"Task {idx+1}/{total_task_number} done.", end="\r")
    return results


def execute_task_in_sequential(
    task: ty.Callable, ranges: ty.List[RangeLike], verbose: bool = True,
) -> ty.List[ty.Any]:
    total_task_number = 1
    for range in ranges:
        total_task_number *= len(range)
    results = []
    for idx, result in enumerate(map(task, product(*ranges))):
        results.append(result)
        if verbose:
            logger.info(f"Task {idx+1}/{total_task_number} done.")
    return results
